# FindClouds.py
# 
# Copyright (C) 2010 Hurng-Chun Lee <hurngchunlee@gmail.com>
# 
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import time
import re

from Ganga.Utility.GridShell import getShell

from GangaAtlas.Lib.AtlasLCGRequirements.AtlasLCGRequirements import _loadCESEInfo, _downloadCESEInfo

from dq2.info import TiersOfATLAS
from dq2.common.DQException import *

def urisplit(uri):
   """
   Basic URI Parser according to STD66 aka RFC3986

   >>> urisplit("scheme://authority/path?query#fragment")
   ('scheme', 'authority', 'path', 'query', 'fragment')

   """
   # regex straight from STD 66 section B
   regex = '^(([^:/?#]+):)?(//([^/?#]*))?([^?#]*)(\?([^#]*))?(#(.*))?'
   p = re.match(regex, uri).groups()
   scheme, authority, path, query, fragment = p[1], p[3], p[4], p[6], p[8]
   #if not path: path = None
   return (scheme, authority, path, query, fragment)

def get_srm_endpoint(site):
    '''
    Gets the SRM endpoint of a site registered in TiersOfATLAS.
    '''

    srm_endpoint_info = {'token':None, 'endpt':None}
    re_srm2 = re.compile('^token:(.*):(srm:\/\/.*)\s*$')

    tmp = TiersOfATLAS.getSiteProperty(site,'srm')
    if tmp:
        srm_endpoint_info['endpt'] = tmp

    mat = re_srm2.match(tmp)
    if mat:
        srm_endpoint_info['token'] = mat.group(1)
        srm_endpoint_info['endpt'] = mat.group(2)

    return srm_endpoint_info

def get_srm_host(site):
    '''
    Gets the SRM hostname of the given site.
    '''
    srm_endpoint_info = get_srm_endpoint(site)

    authority = urisplit(srm_endpoint_info['endpt'])[1]

    return authority.split(':')[0]

def findClouds(job, use_blacklist=True):
    """matching atlas_rtags availability with dataset location, returnning a list of possible clouds"""

    app = job.application


    ## please note the logic to compose ATLAS releaset tag only applicable for
    ## release >= 15
    atlas_rtag        = ''
    atlas_rtag_suffix = ''
    atlas_rtag_prefix = 'VO-atlas-offline'

    if app.atlas_production:

        atlas_project = app.atlas_project

        if not atlas_project:
            atlas_project = 'AtlasProduction'

        if atlas_project in ['AtlasPoint1', 'AtlasTier0', 'AtlasProduction']:
            atlas_rtag_prefix = 'VO-atlas-%s' % re.sub('^Atlas', '', app.atlas_project).lower()
        else:
            atlas_rtag_prefix = 'VO-atlas-%s' % app.atlas_project.lower()

        atlas_rtag_suffix = '%s-%s' % (app.atlas_production, app.atlas_cmtconfig)
    else:
        atlas_rtag_suffix = '%s-%s' % (app.atlas_release, app.atlas_cmtconfig)

    atlas_rtag = '%s-%s' % (atlas_rtag_prefix, atlas_rtag_suffix)

    print 'Atlas release Tag: %s' % atlas_rtag

    bl_locs = []

    ## retrieve blacklisted site
    if use_blacklist:
        ceseinfo = _loadCESEInfo()
        if not ceseinfo or ( time.time() - ceseinfo['time'] > 3600 ):
            ceseinfo_new = _downloadCESEInfo()
            if ceseinfo_new:
                ceseinfo = ceseinfo_new

        if ceseinfo:
            bl_locs = ceseinfo['blacklist']

    ## getting the SE name close to the CE where the required ATLAS release tag is deployed
    cmd = 'lcg-info --vo atlas --list-ce --query \'Tag=%s\' --attrs CloseSE --sed' % atlas_rtag

    s = getShell('GLITE')
    rc, out, m = s.cmd1(cmd,allowed_exit=[0,255])

    se_list = []
    if rc == 0:
        for l in out.strip().split('\n'):
            l.strip()
            ce,ses = l.split('%')
            se_list += ses.split('&')

    ##print se_list

    ## getting the dataset locations and matching it with release locations
    good_locs_map = {}
    for dataset in job.inputdata.dataset:

        good_locs_map[dataset] = {}

        ds = DQ2Dataset()
        ds.dataset = [ dataset ]
        locs = ds.get_locations()

        good_locs = []
        for loc in locs:
            if loc not in bl_locs:
                se = get_srm_host(loc)
                if se in se_list:
                    good_locs.append(loc)

        ##print good_locs

        ## convert good_locs into Cloud info
        dbcloud = {}
        for dbc,cloud in TiersOfATLAS.ToACache.dbcloud.items():
            dbcloud[cloud] = dbc

        good_clouds = {}
        for loc in good_locs:
            c = dbcloud[ TiersOfATLAS.whichCloud(loc) ]

            if not good_clouds.has_key( c ):
                good_clouds[c] = []

            good_clouds[c].append(loc)

    ##print good_clouds

        good_locs_map[dataset] = good_clouds

    return good_locs_map
    
